# python学习 至此开始
# 学习日期: 2023/8/28 19:00
# 学习者: 任磊


# 导入所需的库文件
import sys
import time
from socket import *
from gmssl import sm2, sm3, sm4, func
from base64 import b64decode
import ClientVerify


# 配置加解密环境
# sm2 配置
# 自己的公私钥
my_prk = "036e29c4ce1f17b5fd35c88e81793bc9de53f46b3766b779b297e062af958405"  # 私钥
my_puk = "8d22779ef058a61365d8a427dfd4df5f661535cb515a05567436ab0bd7ff7803a9e31c90630d9221960e614a4228a0a5162bde4072ef3aa9a1ba6fe74d240577"  # 公钥


def my_s2_decry_function(text):  # 定义自己的私钥解密函数
    my_s2 = sm2.CryptSM2(my_prk, my_puk)  # 实例化一个自己的 sm2 对象
    clear_text = b64decode(text.encode())  # 通过base64解码成二进制bytes
    clear_text = my_s2.decrypt(clear_text).decode(encoding="utf-8")
    return clear_text  # 返回密文 byte格式


# sm3 配置
def s3_function(text):  # 定义 sm3 哈希函数
    hash_text = text.encode('UTF-8')  # 将消息进行编码
    hash_text = sm3.sm3_hash(func.bytes_to_list(hash_text))  # 将消息进行哈希
    return hash_text  # 返回哈希值  str 型


# sm4 配置
def s4_encry_function(key, text):  # sm4 加密函数
    s4 = sm4.CryptSM4()  # 实例化一个对象
    s4.set_key(key.encode(), sm4.SM4_ENCRYPT)  # 设置加密密钥
    text = str(text)  # 将明文转为字符串型
    secret_text = s4.crypt_ecb(text.encode())  # 将明文编码后，使用密钥进行加密
    return secret_text  # 返回密文  byte 型


def s4_decry_function(key, text):  # sm4 解密函数，此处得确保录入的密文为 byte 型
    s4 = sm4.CryptSM4()  # 实例化一个对象
    s4.set_key(key.encode(), sm4.SM4_DECRYPT)  # 设置解密密钥
    clear_text = s4.crypt_ecb(text)  # 对 byte 型密文解密
    return clear_text  # 返回明文  byte 型


def get_host_ip():  # 获取本机局域网 ip 函数
    s = socket(AF_INET, SOCK_DGRAM)
    try:
        s.connect(('8.8.8.8', 80))
        ip = s.getsockname()[0]
    finally:
        s.close()
    return ip


IP = get_host_ip()  # IPV4
SEVER_PORT = 9000  # 端口
Max_data = 2048  # 缓冲区大小
dataSocket = socket(AF_INET, SOCK_STREAM)
print('通信环境配置完毕，正在尝试连接...')
dataSocket.connect((IP, SEVER_PORT))  # 连接至目标
print("成功连接，正在接收消息...")
data = dataSocket.recv(Max_data)
other_puk = data.decode()
print("已接收到对方公钥:{}".format(other_puk))
data = dataSocket.recv(Max_data)
other_puk_sign = data.decode()
print("已接收到已签名公钥数据:{}".format(other_puk_sign))
print("正在使用本地公钥进行验签")
if ClientVerify.system_verify_message(other_puk_sign, other_puk):
    print("签名验证通过")
else:
    print("\033[91m签名验证未通过，存在风险！\033[0m")
try:
    data1 = dataSocket.recv(Max_data)  # 接收数字信封
    connect_key = my_s2_decry_function(data1.decode())  # 使用自己的私钥解密
    print(f"已成功接收并解密到通信密钥:{connect_key}")
    data2 = dataSocket.recv(Max_data)  # 接收密文
    target_key = s4_decry_function(connect_key, data2)  # 使用通信密钥解密密文
    target_key = target_key.decode()  # 解码
    print(f"已成功接收并解密到目标密钥:{target_key}")
    print("正在验证目标密钥...")
    # 验证签名
    rec_key_hash = s3_function(target_key)  # 计算接收到的目标密钥的哈希值
    hash_data = dataSocket.recv(Max_data)   # 接收对方发送的签名
    print("目标密钥哈希后为:{}".format(rec_key_hash))
    print("已接收到签名:{}".format(hash_data))
    print("正在进行验签")
    if other_s2.verify(hash_data, rec_key_hash.encode()):  # 验证签名
        print("验证签名通过")
    else:
        print("验证签名未通过")
finally:
    dataSocket.close()  # 关闭通信
    print('通信结束')
